import os
import dashscope
from dashscope import MultiModalConversation
import json
from textGen3D import TextTo3DGenerator
from qwenTempUpload import upload_file_and_get_url
import shutil

# 若使用新加坡地域的模型，请取消下列注释
# dashscope.base_http_api_url = "https://dashscope-intl.aliyuncs.com/api/v1"
def extract_first_json(text: str) -> str:
    start = text.find('{')
    if start == -1:
        raise ValueError('未找到 JSON 开始标记')
    depth = 0
    for i in range(start, len(text)):
        ch = text[i]
        if ch == '{':
            depth += 1
        elif ch == '}':
            depth -= 1
            if depth == 0:
                return text[start:i+1]
    raise ValueError('JSON 括号不平衡')


def convert_to_cubes_params(data):
    # 支持 data 为 dict({objects: [...]}) 或直接为 list([...])
    if isinstance(data, list):
        objs = data
    else:
        objs = (data or {}).get('objects', [])
    cubes = []
    for obj in objs:
        # 仅处理字典对象，其它类型跳过
        if not isinstance(obj, dict):
            continue
        size = obj.get('size', {})
        base = obj.get('base_center', {})
        rot = obj.get('rotation_deg', {})
        rz = obj.get('rotation_z', None)
        # rotation_z 既可能是数值，也可能误写为带有 yaw/pitch/roll 的字典，做兼容处理
        if isinstance(rz, dict):
            rz_value = rz.get('yaw', 0)
        else:
            try:
                rz_value = float(rz) if rz is not None else None
            except Exception:
                rz_value = None
        # 兼容 size 为 list/tuple 的情况
        if isinstance(size, (list, tuple)) and len(size) == 3:
            dx, dy, dz = size
        elif isinstance(size, dict):
            dx = size.get('dx', size.get('width', 1))
            dy = size.get('dy', size.get('depth', 1))
            dz = size.get('dz', size.get('height', 1))
        else:
            dx, dy, dz = 1, 1, 1
        # 兼容 base_center 为 list/tuple 的情况
        if isinstance(base, (list, tuple)) and len(base) == 3:
            bx, by, bz = base
        elif isinstance(base, dict):
            bx = base.get('x', 0)
            by = base.get('y', 0)
            bz = base.get('z', 0)
        else:
            bx, by, bz = 0, 0, 0
        cubes.append({
            'size': [dx, dy, dz],
            'position': [bx, by, bz],
            'color': obj.get('color', '#4C9AFE'),
            'rotation_x': rot.get('roll', 0),   # roll -> x
            'rotation_y': rot.get('pitch', 0),  # pitch -> y
            'rotation_z': rz_value if rz_value is not None else rot.get('yaw', 0),    # prefer rotation_z, fallback to yaw
        })
    return cubes

def sanitize_to_z_only(data):
    # 兼容模型返回为 list 的情况：统一包装为 {objects: list}
    if isinstance(data, list):
        data = { 'objects': data }
    objs = (data or {}).get('objects', [])
    for obj in objs:
        if not isinstance(obj, dict):
            # 跳过非字典类型的对象
            continue
        # Normalize rotation_z to a numeric value
        if 'rotation_z' in obj:
            rz = obj['rotation_z']
            if isinstance(rz, dict):
                obj['rotation_z'] = rz.get('yaw', 0)
            else:
                try:
                    obj['rotation_z'] = float(rz)
                except Exception:
                    obj['rotation_z'] = 0
        else:
            rot = obj.get('rotation_deg', {})
            obj['rotation_z'] = rot.get('yaw', 0)
        # Remove any non-z rotation fields
        obj.pop('rotation_deg', None)
        obj.pop('rotation_x', None)
        obj.pop('rotation_y', None)
    return data
# 0 30 60
example_url = "oss://dashscope-instant/8737449652d607f55844dc72848f38e6/2025-10-27/e49628d7-1fad-4d84-877d-6b7466698bf6/rotated_y_cuboids_with_plane.png"

target_url = "oss://dashscope-instant/8737449652d607f55844dc72848f38e6/2025-10-27/72ed47d3-5775-4684-9e3d-d78a429d2723/2276.jpg_wh860.jpg"

messages = [
    {
        "role": "user",
        "content": [
            {"image": example_url},
            {"image": target_url},
            {"text": """你是一名建筑识别专家,你的任务是将一个图片中的所有建筑或者类似长方体的形状抽象为3D空间中的长方体，只需要保证建筑的相对位置大概正确和相对比例大概正确即可。我们提供两张图片：\n- 示例图1为第一张给出的图片\n- 待识别图2（本消息中的第二张图片）\n\n示例图由 textGen3D.py 生成，生成逻辑要点：\n1) 坐标系右手系，z 轴向上，z=0 为地面；仅显示 z>=0 空间。\n2) 每个长方体以 base_center 表示其在地面的中心位置；size={dx,dy,dz} 为沿 x/y/z 的尺寸。\n3) 所有长方体仅沿 z 轴旋转，必须包含数值型 rotation_z 字段（单位：度，+z 轴方向逆时针为正）；不要返回 yaw/pitch/roll 或 rotation_deg。\n\n示例图（图1）中的三个长方体数据为：\n- 红色长方体：size=[8,4,3]，base_center=(0,0,0)，rotation_z=0°；\n- 绿色长方体：size=[6,3,2]，base_center=(12,0,0)，rotation_z=30°；\n- 蓝色长方体：size=[4,2,1]，base_center=(20,0,0)，rotation_z=60°。\n\n示例图（图1）中三个长方体的 rotation_z 分别为 0、30、60 度。\n\n示例图（图1）的标准标注如下，后续请严格按相同键名与层级返回：\n{\n  \"version\": \"1.0\",\n  \"coordinate_system\": \"right-handed-z-up\",\n  \"objects\": [\n    {\n      \"type\": \"cuboid\",\n      \"id\": 1,\n      \"base_center\": {\"x\": 0, \"y\": 0, \"z\": 0},\n      \"size\": {\"dx\": 8, \"dy\": 4, \"dz\": 3},\n      \"rotation_z\": 0,\n      \"color\": \"#FF0000\"\n    },\n    {\n      \"type\": \"cuboid\",\n      \"id\": 2,\n      \"base_center\": {\"x\": 12, \"y\": 0, \"z\": 0},\n      \"size\": {\"dx\": 6, \"dy\": 3, \"dz\": 2},\n      \"rotation_z\": 30,\n      \"color\": \"#00FF00\"\n    },\n    {\n      \"type\": \"cuboid\",\n      \"id\": 3,\n      \"base_center\": {\"x\": 20, \"y\": 0, \"z\": 0},\n      \"size\": {\"dx\": 4, \"dy\": 2, \"dz\": 1},\n      \"rotation_z\": 60,\n      \"color\": \"#0000FF\"\n    }\n  ]\n}\n\n任务：以图1为示例规范，识别图2中的所有建筑，并仅输出严格符合上述结构的 JSON（不要附加解释）。角度无法判断时统一填 0。"""},
        ],
    }
]

# 简易的输出辅助，能在失败时打印错误信息

def show_response(tag, resp):
    try:
        print(f"{tag}：{resp.output.choices[0].message.content[0]['text']}")
    except Exception as e:
        print(f"{tag} 调用失败: {e}")
        # 常见字段提示
        msg = getattr(resp, 'message', None) or getattr(resp, 'error', None) or getattr(resp, 'code', None)
        if msg:
            print("原始响应：", msg)
        else:
            print("原始响应对象：", resp)

response = MultiModalConversation.call(
    # 若没有配置环境变量，请用百炼API Key将下行替换为：api_key="sk-xxx",
    # 新加坡和北京地域的API Key不同。获取API Key：https://help.aliyun.com/zh/model-studio/get-api-key
    api_key=os.getenv('DASHSCOPE_API_KEY'),
    model='qwen3-vl-plus',   # 模型名称按需更换
    messages=messages)
show_response("模型第一轮输出", response)

# 根据第一轮结构化数据生成 3D 立方体
try:
    first_text = response.output.choices[0].message.content[0]['text']
    data = json.loads(extract_first_json(first_text))
    data = sanitize_to_z_only(data)
    print('规范化后的 JSON：')
    print(json.dumps(data, ensure_ascii=False, indent=2))
    cubes_params = convert_to_cubes_params(data)
    generator = TextTo3DGenerator(output_dir='output/text3d')
    first_round_img_path = generator.generate_multi_cubes(cubes_params, 'from_qwen_first_round', show_z_plane=True)
    print(f'已根据第一轮数据生成 3D 模型: {first_round_img_path}')

    # 上传第一轮生成的图片到临时OSS
    api_key = os.getenv('DASHSCOPE_API_KEY')
    try:
        first_round_oss_url = upload_file_and_get_url(api_key, 'qwen3-vl-plus', first_round_img_path)
        print(f'第一轮生成图片已上传到OSS: {first_round_oss_url}')
    except Exception as ue:
        first_round_oss_url = None
        print('上传第一轮图片到OSS失败:', ue)

    # 迭代优化过程：将每次新调整的3D数据与生成图片保存到 latest_gen 前缀，并再次与 target_url 对比
    latest_json = data
    latest_img_path = first_round_img_path
    latest_img_oss = first_round_oss_url
    latest_json_path = os.path.join('output', 'text3d', 'latest_gen.json')
    latest_png_path = os.path.join('output', 'text3d', 'latest_gen.png')
    try:
        with open(latest_json_path, 'w', encoding='utf-8') as f:
            f.write(json.dumps(latest_json, ensure_ascii=False, indent=2))
        shutil.copyfile(latest_img_path, latest_png_path)
    except Exception:
        pass

    num_iterations = int(os.getenv('QWEN_ADJUST_ITERS', '2'))  # 默认进行2次迭代
    for i in range(5):
        if not latest_img_oss:
            # 如果当前图未能上传 OSS，则尝试重新上传
            try:
                latest_img_oss = upload_file_and_get_url(api_key, 'qwen3-vl-plus', latest_img_path)
                print(f'第{i+1}轮：latest_gen 图片已上传到OSS: {latest_img_oss}')
            except Exception as ue:
                print(f'第{i+1}轮：上传 latest_gen 图片到OSS失败，终止迭代:', ue)
                break

        adjust_prompt = (
            f'第{i+1}轮调整：图A为当前 latest_gen，图B为目标图。'
            '结合下面给出的 latest_gen 结构化数据，返回让图A更接近图B的调整后的 JSON。'
            '要求：仅沿 z 轴旋转，必须使用数值型 rotation_z（单位：度）；'
            '如需调整，请直接更新 objects 中的 size/base_center/rotation_z/color。\n\n'
            'latest_gen 结构化数据：\n' + json.dumps(latest_json, ensure_ascii=False, indent=2)
        )
        enable_thinking = True
        messages_iter = [{
            'role': 'user',
            'content': [
                {'image': latest_img_oss},
                {'image': target_url},
                {'text': adjust_prompt},
            ],
        }]

        response_iter = MultiModalConversation.call(
            api_key=api_key,
            model='qwen3-vl-plus',
            messages=messages_iter,
            enable_thinking=enable_thinking,
            # thinking_budget 参数设置最大推理过程 Token 数，仅对qwen-vl-plus、 qwen3-vl-plus-2025-09-23，qwen3-vl-235b-a22b-thinking
            thinking_budget=50,
        )
        show_response(f'模型第{i+1}轮输出', response_iter)

        try:
            iter_text = response_iter.output.choices[0].message.content[0]['text']
            iter_json = json.loads(extract_first_json(iter_text))
            print(f'第{i+1}轮建议 JSON：')
            print(json.dumps(iter_json, ensure_ascii=False, indent=2))
            # 保存该轮建议
            sug_path = os.path.join('output', 'text3d', f'latest_gen_round{i+1}_suggestion.json')
            with open(sug_path, 'w', encoding='utf-8') as f:
                f.write(json.dumps(iter_json, ensure_ascii=False, indent=2))
            # 生成该轮 3D 模型
            iter_json_norm = sanitize_to_z_only(iter_json)
            cubes_params_iter = convert_to_cubes_params(iter_json_norm)
            latest_img_path = generator.generate_multi_cubes(cubes_params_iter, f'latest_gen_round{i+1}', show_z_plane=True)
            print(f'已生成第{i+1}轮 3D 模型: {latest_img_path}')
            # 更新 latest_gen 快照
            latest_json = iter_json_norm
            with open(latest_json_path, 'w', encoding='utf-8') as f:
                f.write(json.dumps(latest_json, ensure_ascii=False, indent=2))
            shutil.copyfile(latest_img_path, latest_png_path)
            # 上传用于下一轮对比
            try:
                latest_img_oss = upload_file_and_get_url(api_key, 'qwen3-vl-plus', latest_img_path)
            except Exception as ue:
                latest_img_oss = None
                print(f'第{i+1}轮：上传生成图到OSS失败，下一轮将重试:', ue)
        except Exception as pe:
            print(f'第{i+1}轮 JSON 解析失败，终止迭代:', pe)
            break
except Exception as e:
    print('第一轮数据转 3D 失败:', e)